package http

import (
	"bytes"
	"crypto/tls"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
	"kscan/pkg/poc"
)

// Client HTTP客户端
type Client struct {
	client  *http.Client
	timeout time.Duration
}

// NewClient 创建HTTP客户端
func NewClient(timeout time.Duration, proxyURL string) (*Client, error) {
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true, // 忽略证书验证
		},
		MaxIdleConns:        100,
		MaxIdleConnsPerHost: 10,
		IdleConnTimeout:     90 * time.Second,
	}

	// 设置代理
	if proxyURL != "" {
		proxy, err := url.Parse(proxyURL)
		if err != nil {
			return nil, err
		}
		transport.Proxy = http.ProxyURL(proxy)
	}

	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse // 默认不跟随重定向
		},
	}

	return &Client{
		client:  client,
		timeout: timeout,
	}, nil
}

// DoRequest 执行HTTP请求
func (c *Client) DoRequest(target string, pocReq poc.Request) (*poc.Response, error) {
	// 构建完整URL
	targetURL := target
	if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
		targetURL = "http://" + target
	}
	
	// 添加路径
	if pocReq.Path != "" {
		if strings.HasPrefix(pocReq.Path, "/") {
			targetURL = strings.TrimSuffix(targetURL, "/") + pocReq.Path
		} else {
			targetURL = strings.TrimSuffix(targetURL, "/") + "/" + pocReq.Path
		}
	}

	// 创建请求
	var bodyReader io.Reader
	if pocReq.Body != "" {
		bodyReader = bytes.NewBufferString(pocReq.Body)
	}

	req, err := http.NewRequest(pocReq.Method, targetURL, bodyReader)
	if err != nil {
		return nil, err
	}

	// 设置请求头
	req.Header.Set("User-Agent", "KSCAN/1.0")
	for key, value := range pocReq.Headers {
		req.Header.Set(key, value)
	}

	// 记录开始时间
	startTime := time.Now()

	// 执行请求
	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// 计算响应时间
	elapsed := time.Since(startTime).Milliseconds()

	// 读取响应体
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 构建响应对象
	response := &poc.Response{
		StatusCode: resp.StatusCode,
		Headers:    resp.Header,
		Body:       string(body),
		Time:       elapsed,
	}

	return response, nil
}

// DoRequestWithRedirect 执行HTTP请求（支持重定向）
func (c *Client) DoRequestWithRedirect(target string, pocReq poc.Request) (*poc.Response, error) {
	// 创建临时客户端支持重定向
	tempClient := &http.Client{
		Transport: c.client.Transport,
		Timeout:   c.timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= pocReq.MaxRedirects {
				return http.ErrUseLastResponse
			}
			return nil
		},
	}

	// 构建完整URL
	targetURL := target
	if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
		targetURL = "http://" + target
	}
	
	if pocReq.Path != "" {
		if strings.HasPrefix(pocReq.Path, "/") {
			targetURL = strings.TrimSuffix(targetURL, "/") + pocReq.Path
		} else {
			targetURL = strings.TrimSuffix(targetURL, "/") + "/" + pocReq.Path
		}
	}

	// 创建请求
	var bodyReader io.Reader
	if pocReq.Body != "" {
		bodyReader = bytes.NewBufferString(pocReq.Body)
	}

	req, err := http.NewRequest(pocReq.Method, targetURL, bodyReader)
	if err != nil {
		return nil, err
	}

	// 设置请求头
	req.Header.Set("User-Agent", "KSCAN/1.0")
	for key, value := range pocReq.Headers {
		req.Header.Set(key, value)
	}

	// 记录开始时间
	startTime := time.Now()

	// 执行请求
	resp, err := tempClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// 计算响应时间
	elapsed := time.Since(startTime).Milliseconds()

	// 读取响应体
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 构建响应对象
	response := &poc.Response{
		StatusCode: resp.StatusCode,
		Headers:    resp.Header,
		Body:       string(body),
		Time:       elapsed,
	}

	return response, nil
}